import csv
import matplotlib.pyplot as plt
from scipy import signal
from sklearn import preprocessing
import numpy as np 
import os
import scipy.io as scio
import torch 
import scipy

def read_file_name(path):
    files = os.listdir(path)
    s = []
    for file in files:
        if not os.path.isdir(file):
            s.append(file)
    return s

def preprocess_raweeg(path_name, debug = False):
    f = open('./info.csv', 'r')
    reader = csv.reader(f)
    data_csv = []
    for item in reader:
        data_csv.append(item)
    f.close()
    file_names = read_file_name(path_name)

    target_data = []
    target_label = []
    index=0
    for fn in file_names:
        tmp_s = fn.split('.')[0].split('_')
        if tmp_s[0] == '':
            continue
        if debug:
            print(tmp_s)
        subject_index = int(tmp_s[0])
        exp_index = int(tmp_s[1])
        start_line = (subject_index - 1) * 6 + (exp_index - 1) * 2
        start_line = int(start_line)
        if 1 <= subject_index <= 15:
            start_point_list = [int(x) for x in data_csv[start_line][:15]]
            end_point_list = [int(x) for x in data_csv[start_line + 1][:15]]
            labels = [1, 0, -1, -1, 0, 1, -1, 0, 1, 1, 0, -1, 0, 1, -1]
        elif 22 <= subject_index <= 29:
            start_point_list = [int(x) for x in data_csv[start_line][:21]]
            end_point_list = [int(x) for x in data_csv[start_line + 1][:21]]
            labels = [1, -1, 0, -1, 1, 0, -1, 0, 1, -1, 0, 1, 1, 0, -1, -1, 0, 1, -1, 0, 1]
        else:
            print('error: subject index illegal')
            return
        if debug:
            print(path_name+fn)
            print(start_point_list)
            print(end_point_list)
            print(labels)
        index = index+1
        print(index)
        data = scio.loadmat(path_name+fn)['raweeg']

        for start_second, end_second, label in zip(start_point_list, end_point_list, labels):
            for i in range(start_second, end_second):
                tmp_data = data[:, i * 1000:(i + 1) * 1000]
            
                transpose_data = np.transpose(tmp_data)
                
                target_data.append(transpose_data)
                target_label.append(label+1) # -1 0 1 to 0 1 2
    return target_data, target_label

def data_process(path_name, start=0, end=-1, resamples=250, is_test=False):
    src, tgt = preprocess_raweeg(path_name)
    src = src[start: end]
    tgt = tgt[start: end]
    
    # src = torch.tensor(src, dtype=torch.float32)
    # tgt = torch.tensor(tgt)
    # print(src.shape)
    # print(tgt.shape)
    if is_test is False:
        torch.save(src, "./data/train_src" + str(start) + "_" + str(end) + ".pkl")
        torch.save(tgt, "./data/train_tgt" + str(start) + "_" + str(end) + ".pkl")
    else:
        torch.save(src, "./data/test_src" + str(start) + "_" + str(end) + ".pkl")
        torch.save(tgt, "./data/test_tgt" + str(start) + "_" + str(end) + ".pkl")



if __name__ == "__main__":
    src, tgt = preprocess_raweeg("./TrainData/")
    # print(type(src))
    # print(np.mean(src))
    # src = torch.tensor(src)
    # print(src.shape)
    # src = src[:100]
    # tgt = torch.tensor(tgt)
    # tgt = tgt[:100]
    # print(len(tgt))
    
    torch.save(tgt, "./data/train_tgt.pkl")
    # src = torch.load("./src.pkl")
    # tgt = torch.load("./tgt.pkl")
    # Y = scipy.signal.resample(src, 250, axis=1)
    src = torch.tensor(src, dtype=torch.float32)
    torch.save(src, "./data/train_src.pkl")
    # print(Y)
    # print(Y.shape)
    # print(src[10:])
    